from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict
from planner import Planner
from writer import Writer
from typing import Annotated
from langgraph.types import Send
from operator import add


class Outline(TypedDict):
    title: str
    sub_titles: list[str]


class SchemeState(TypedDict):
    project_name: str
    abstract: str
    outlines: list[Outline]
    contents: Annotated[list[str], add]
    article: str


class ContentState(TypedDict):
    project_name: str
    abstract: str
    title: str
    sub_title: str


class Designer:

    def __init__(self):
        _llm = ChatOpenAI(
            api_key="ollama",
            model="qwen2.5:7b",
            base_url="http://192.168.10.13:60001/v1",
            temperature=0.7,
        )

        self._planner = Planner(_llm)
        self._writer = Writer(_llm)

        self._graph = self._init_graph()

    def _init_graph(self):
        _builder = StateGraph(SchemeState)
        _builder.add_node("_planner_node", self._planner_node)
        _builder.add_node("_writer_node", self._writer_node)
        _builder.add_node("_gather_node", self._gather_node)

        _builder.add_edge(START, "_planner_node")
        _builder.add_conditional_edges("_planner_node", self._dispatch_edge)
        _builder.add_edge("_writer_node", "_gather_node")
        _builder.add_edge("_gather_node", END)

        _graph = _builder.compile()
        return _graph

    def _planner_node(self, state):
        while True:
            try:
                return self._planner(state)
            except Exception as e:
                print(e)

    def _writer_node(self, state: ContentState):
        while True:
            try:
                _rt = self._writer(state)
                return {"contents": [_rt["content"]]}
            except Exception as e:
                print(e)

    def _gather_node(self, state):
        _article = f"{state['project_name']}\n\n"
        _k = 0
        for _i, _outline in enumerate(state["_outlines"]):
            _article += f"{_i+1}.{_outline['title']}\n\n"
            for _j, _sub_title in enumerate(_outline["sub_titles"]):
                _article += f"{_i+1}.{_j+1}.{_sub_title}\n\n"
                _article += f"{state['contents'][_k]}\n\n\n\n"
                _k += 1

    def _dispatch_edge(self, state):
        _project_name = state["project_name"]
        _abstract = state["abstract"]
        _outlines = state["outlines"]

        ret = []
        for _outline in _outlines:
            _title = _outline["title"]
            print(_outline)
            for _sub_title in _outline["sub_titles"]:
                ret.append(
                    Send(
                        "_writer_node",
                        {
                            "project_name": _project_name,
                            "abstract": _abstract,
                            "title": _title,
                            "sub_title": _sub_title,
                        },
                    )
                )
            return ret

    def __call__(self, project_name):
        return self._graph.invoke({"project_name": project_name})


if __name__ == "__main__":
    designer = Designer()
    rt = designer("基于多模态大模型驱动的数字人")
    article = rt["article"]
    with open("article.txt", "w+") as f:
        f.write()
